import pandas as pd
import xlsxwriter as xl
from xlsxwriter.utility import xl_range
from utils import *

pd.set_option("display.unicode.east_asian_width", True)     # print时对齐中文字符
# pd.set_option('display.max_columns', None) # 设置最大显示列数
# pd.set_option('display.max_rows', None) # 设置最大显示列数
# pd.set_option('display.width', None) # 设置宽度

# 修改为只使用年份列表，不再按批次区分
years = [2024, 2023, 2022, 2021, 2020, 2019, 2018, 2017, 2016, 2015, 2014]
subjects = {'science': '理工类', 'arts': '文史类'}

def generate_data(years):
    rank_history = {
        'science': pd.DataFrame(),
        'arts': pd.DataFrame()
    }
    
    # 读取所有投档线数据
    admission_all = pd.read_csv(admission_all_in_one_path(), dtype=ADMISSION_DTYPES)

    for subject, subject_zh in subjects.items():
        # 设置rank_history的列名、索引等
        # 索引： ['院校代码', '专业组']， 列：[(year, '名次'), (year, '投档线'), ...]
        rank_history[subject] = pd.DataFrame(
            columns = [(year, c) for year in years for c in ['名次', '投档线']],
            index=pd.MultiIndex.from_tuples([], names=['院校代码', '专业组'])
        )
        # 找到所有存在的院校代码和专业组
        admission_subject = admission_all[admission_all['类别'] == subject].copy()
        admission_subject = admission_subject[['院校代码', '专业组']].drop_duplicates()
        # 丢弃NaN
        admission_subject = admission_subject[admission_subject['院校代码'].notna() & admission_subject['专业组'].notna()]
        # 添加所有行到rank_history的索引中
        rank_history[subject] = rank_history[subject].reindex(admission_subject.set_index(['院校代码', '专业组']).index)
        print(f'初始化rank_history({subject}): \n', rank_history[subject])

        # 遍历每个年份，处理数据
        for year in years:
            # 过滤出当前年份和科目的数据
            admission = admission_all[
                (admission_all['年份'] == year) & 
                (admission_all['类别'] == subject)
            ].copy()
            
            # 读取一分一档表获取排名信息
            score_rank_path = score_rank_data_path(year, subject)
            score_rank = pd.read_csv(score_rank_path, dtype=SCORE_RANK_DTYPES)
            score_rank.set_index('分值', inplace=True)
            # 为了方便查询，将不存在的分数值补充完整
            # 在某些年份，不公布前10名时，前10名的名次均显示为0
            score_rank = score_rank.reindex(pd.RangeIndex(MAX_SCORE, MIN_SCORE, -1))
            score_rank['人数'].fillna(0, inplace=True)
            score_rank['名次'] = score_rank['名次'].ffill().fillna(1)
            score_rank['累计人数'] = score_rank['累计人数'].ffill().fillna(1)

            # 对 admission 中的每一行，将数据填入 rank_history 中；
            # 有两种情况，专业组为 NaN，则需要对该校每个存在的专业组都填入对应的名次和投档线；
            # 否则，只填入对应专业组的名次和投档线。
            for _, row in admission.iterrows():
                school_code = row['院校代码']
                major_group = row['专业组']
                score = row['投档线']
                # 使用score_rank获取对应的名次
                if pd.isna(score):
                    # 如果投档线为 NaN，则跳过该行
                    continue
                rank = score_rank.loc[score, '名次']
                # print(f'学校:{school_code}；专业组:{major_group}；投档线:{score}；名次：{rank}')
                if pd.isna(major_group):
                    # 如果专业组为 NaN，则对该校每个存在的专业组都填入对应的名次和投档线
                    for group in admission_subject[admission_subject['院校代码'] == school_code]['专业组']:
                        # 只有当当前值为NaN或者当前名次比新名次更大时才更新
                        current_rank = rank_history[subject].at[(school_code, group), (year, '名次')]
                        if pd.isna(current_rank) or current_rank > rank:
                            rank_history[subject].at[(school_code, group), (year, '名次')] = rank
                            rank_history[subject].at[(school_code, group), (year, '投档线')] = score
                else:
                    # 否则，只填入对应专业组的名次和投档线
                    current_rank = rank_history[subject].at[(school_code, major_group), (year, '名次')]
                    if pd.isna(current_rank) or current_rank > rank:
                        rank_history[subject].at[(school_code, major_group), (year, '名次')] = rank
                        rank_history[subject].at[(school_code, major_group), (year, '投档线')] = score
                


    # print("score_rank:", score_rank, sep='\n')
    #print("rank_history:", rank_history['science'], sep='\n')
    # exit(0)

    # 后处理
    # 获取所有学校的名称
    school_names = admission_all.copy()[['院校代码', '院校名称']].drop_duplicates()
    school_names.set_index('院校代码', inplace=True)
    school_names = school_names['院校名称']
    school_names = school_names.to_dict()
    
    ORDER_DATA_TYPE = {'院校代码': 0, '院校名称': 1, '专业组': 2, '名次': 3, '投档线': 4}
    for subject, subject_zh in subjects.items():
        # 使用向量化操作生成院校名称列
        # 提取院校代码并映射到院校名称
        school_codes = rank_history[subject].index.get_level_values('院校代码')
        school_name_col = []
        for code in school_codes:
            school_name_col.append(school_names[code])
        
        # print(f'院校名称列({subject}): \n', len(school_name_col))
        # print(f'{rank_history[subject].shape}')

        print(rank_history[subject])

        # 插入院校名称列
        rank_history[subject].insert(0, '院校名称', school_name_col)
        print(f'插入院校名称({subject}): \n', rank_history[subject])

        # 将目前作为索引的院校代码和专业组列转换为普通列
        rank_history[subject].reset_index(inplace=True)
        print(f'转换索引({subject}): \n', rank_history[subject])

        # 将columns转为两层索引，方便归类（前三列院校代码、名称、专业组年份设置为最新一年，方便排序）
        columns = rank_history[subject].columns
        columns = columns.tolist()
        columns[0] = (max(years)+1, columns[0])
        columns[1] = (max(years)+1, columns[1])
        columns[2] = (max(years)+1, columns[2])
        columns = pd.DataFrame(columns, columns={'年度': 'Int64', '数据': 'object'})
        rank_history[subject].columns = pd.MultiIndex.from_frame(columns[['数据', '年度']])

        # 列按照[院校名称, 历年名次, 历年投档线]的顺序排序
        rank_history[subject].sort_index(axis=1, key=lambda x: x.map(ORDER_DATA_TYPE), inplace=True)
        # 行按照名次升序排序
        rank_history[subject].sort_values(axis=0, by=('名次', years[0]), inplace=True)

        print(f'排序({subject}): \n', rank_history[subject])
    return rank_history


# 将数据写入excel表
def build_excel_table(rank_history, title, file):
    os.makedirs(os.path.dirname(file), exist_ok=True)
    writer = pd.ExcelWriter(file, engine='xlsxwriter')
    for subject, subject_zh in subjects.items():
        print(subject, subject_zh)
        # 准备数据
        table = rank_history[subject][['院校代码', '院校名称', '专业组', '名次']]

        # 院校代码 和 院校名称两列名字单独添加
        columns = ['院校代码', '院校名称', '专业组']
        # 剩余列名为数据+年份（如：名次2020）
        columns.extend([''.join(map(str, col)) for col in table.columns[3:]])
        table.columns = columns

        sheet_name = f'历史名次（{subject_zh}）'
        _, _, nrow, ncol = write_to_excel_table(table, writer, sheet_name, title + f'（{subject_zh}）')
        workbook = writer.book
        worksheet = writer.sheets[sheet_name]

        # 添加迷你图趋势线
        sparkline_options = {
            'markers': True,
            'reverse': True
        }
        SPARKLINE_START_COL = 3         # 从第4列开始添加迷你图（索引从0开始）
        for i in range(table.shape[0]):
            sparkline_options['range'] = xl_range(i+2, SPARKLINE_START_COL, i+2, ncol)
            worksheet.add_sparkline(i+2, ncol+1, sparkline_options)
        
        # 添加说明
        red = workbook.add_format({'color': '#C00000', 'bold': 1})
        blue = workbook.add_format({'color': '#16365C', 'bold': 1})
        cell_format = workbook.add_format({'align': 'left',
                                    'valign': 'top',
                                    'text_wrap': True})
        
        INFO_ROW_START = 2
        INFO_ROW_END = 18
        IMAGE_ROW_START = 19
        IMAGE_ROW_END = IMAGE_ROW_START + 12
        worksheet.merge_range(INFO_ROW_START, ncol+2, INFO_ROW_END, ncol+2, '', cell_format)
        worksheet.set_column_pixels(ncol+2, ncol+2, 220)
        worksheet.write_rich_string(2, ncol+2, 
            blue, '趋势图说明：\n',
            '按时间顺序', red, '从左至右', '绘制，曲线',
            red, '越低', '代表排名越高，', red, '越难录取', '。\n\n',
            blue, '数据说明：\n',
            '数据采用历年投档线对应名次。空白代表当年该校无人录取。部分年份排名靠前几十位的成绩未公布，名次显示为1。\n',
            red, '2023年及之前', '的投档按学校而不是学校+专业组，因此', red, '数据使用当年学校的最低投档线', '而不是专业组的最低投档线。\n',
            red, '2024年起', '本科普通批不再分一本和二本，因此', red, '2023年及以前，一本、二本均有数据的高校，采用二本分数线作为历史数据', '。\n',
            red, '数据仅供参考，请以考试院官方发布数据为准。', cell_format
        )
        cell_format = workbook.add_format({'align': 'center',
                                    'valign': 'top',
                                    'color': '#16365C',
                                    'bold': 1,
                                    'text_wrap': True})
        worksheet.merge_range(IMAGE_ROW_START, ncol+2, IMAGE_ROW_END, ncol+2, '欢迎报考\n中国科学技术大学', cell_format)
        worksheet.insert_image(IMAGE_ROW_START+2, ncol+2, PATH_USTC_LOGO, {'x_scale': 1, 'y_scale': 5 / 4.52, 'x_offset': 30})

    writer.close()

if __name__ == '__main__':
    file = os.path.join(DIR_RELEASE, f'高校专业组历年录取名次趋势表（{min(years)}年~{max(years)}年）.xlsx')
    rank_history = generate_data(years)
    for subject, subject_zh in subjects.items():
        title = f'{min(years)}年至{max(years)}年各高校录取名次趋势表'
        print(title)
    build_excel_table(rank_history, title, file)